from __future__ import with_statement

from objc import Category
from AppKit import *
from os.path import splitext

_fileRepresentationMapping = {
        '.png': NSPNGFileType,
        '.gif': NSGIFFileType,
        '.jpg': NSJPEGFileType,
        '.jpeg': NSJPEGFileType,
        '.bmp': NSBMPFileType,
        '.tif': NSTIFFFileType,
        '.tiff': NSTIFFFileType,
    }

def _getFileRepresentationType(filepath):
    base, ext = splitext(filepath)
    return _fileRepresentationMapping[ext.lower()]

class NSImage(Category(NSImage)):

    def rect(self):
        return (0,0),self.size()

    @classmethod
    def imageWithFilePath_(cls, filepath):
        return NSImage.alloc().initWithContentsOfFile_(filepath)
        
    @classmethod
    def imageWithSize_(cls, size):
        image = NSImage.alloc().initWithSize_(size)
        image.clear()
        return image
        
    def scale_(self, scalingFactor):
        width, height = self.size()
        newsize = width * scalingFactor, height * scalingFactor
        newimage = NSImage.imageWithSize_(newsize)
        with newimage:
            NSGraphicsContext.currentContext().setImageInterpolation_(NSImageInterpolationHigh)
            self.drawInRect_fromRect_operation_fraction_(((0,0),newsize), ((0,0),(width,height)), NSCompositeCopy, 1.0)
        return newimage

    def writeToFilePath_(self, filepath):
        with self:
            image_rep = NSBitmapImageRep.alloc().initWithFocusedViewRect_(self.rect())
        representation = _getFileRepresentationType(filepath)
        data = image_rep.representationUsingType_properties_(representation, None)
        data.writeToFile_atomically_(filepath, False)
        
    def clear(self):
        self.fillWithColor_(NSColor.whiteColor())

    def fillWithColor_(self, color):
        with self:
            color.set()
            NSBezierPath.fillRect_(self.rect())

    # 
    #  Implement context control for with statement
    #         

    def __enter__(self):
        self.lockFocus()
        return self

    def __exit__(self, type, value, tb):
        self.unlockFocus()
        return tb is None # False will re-raise any exception



    
